package com.zhanghu.backendZray.utils.AI;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.fastjson.JSON;
import io.reactivex.Flowable;

/**
 * 官方文档：https://help.aliyun.com/zh/model-studio/getting-started/first-api-call-to-qwen?spm=a2c4g.11186623.0.0.5403b0a8DO1NYT#a0d9906322g41
 */
public class TongyiUtil {

    public static String call(List<String> msgHistory, String content) throws ApiException, NoApiKeyException, InputRequiredException {
        GenerationResult result = callWithMessage(msgHistory, content);
        return result.getOutput().getChoices().get(0).getMessage().getContent();
    }

    /**
     * 一口气全输出
     *
     * @param msgHistory
     * @param content
     * @return
     * @throws ApiException
     * @throws NoApiKeyException
     * @throws InputRequiredException
     */
    public static GenerationResult callWithMessage(List<String> msgHistory, String content) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        List<Message> requestMsgList = new ArrayList<>();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("You are a helpful assistant.")
                .build();
        requestMsgList.add(systemMsg);
        if (CollUtil.isNotEmpty(msgHistory)) {
            for (String msg : msgHistory) {
                Message userMsg = Message.builder()
                        .role(Role.USER.getValue())
                        .content(msg)
                        .build();
                requestMsgList.add(userMsg);
            }
        }
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(content)
                .build();
        requestMsgList.add(userMsg);
        GenerationParam param = GenerationParam.builder()
//                .model("qwen-turbo")
//                .model("qwen2.5-3b-instruct")
                .model("qwen-turbo-latest")
                .messages(requestMsgList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        return gen.call(param);
    }

    /**
     * 流式输出
     *
     * @param msgHistory
     * @param content
     * @return
     * @throws ApiException
     * @throws NoApiKeyException
     * @throws InputRequiredException
     */
    public static Flowable<GenerationResult> callWithMessageStream(List<String> msgHistory, String content) throws ApiException, NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        List<Message> requestMsgList = new ArrayList<>();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("You are a helpful assistant.")
                .build();
        requestMsgList.add(systemMsg);
        if (CollUtil.isNotEmpty(msgHistory)) {
            for (String msg : msgHistory) {
                Message userMsg = Message.builder()
                        .role(Role.USER.getValue())
                        .content(msg)
                        .build();
                requestMsgList.add(userMsg);
            }
        }
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(content)
                .build();
        requestMsgList.add(userMsg);
        GenerationParam param = GenerationParam.builder()
//                .model("qwen-turbo")
//                .model("qwen2.5-3b-instruct")
                .model("qwen-turbo-latest")
                .messages(requestMsgList)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .build();
        return gen.streamCall(param);
    }

    public static void main1(String[] args) {
        try {
            GenerationResult result = callWithMessage(Collections.emptyList(), "你是谁");
            System.out.println(result.getOutput().getChoices().get(0).getMessage().getContent());
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            // 使用日志框架记录异常信息
            System.err.println("An error occurred while calling the generation service: " + e.getMessage());
        }
        System.exit(0);
    }

    private static void handleGenerationResult(GenerationResult message) {
        for (GenerationOutput.Choice choice : message.getOutput().getChoices()) {
            String content = choice.getMessage().getContent();
            System.err.println(content);
        }
//        System.err.println(JSON.toJSONString(message));
    }

    public static Flowable<GenerationResult> streamCallWithMessage(String question)
            throws NoApiKeyException, ApiException, InputRequiredException {
        Generation gen = new Generation();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(question).build();
        GenerationParam param = buildGenerationParam(userMsg);
        Flowable<GenerationResult> result = gen.streamCall(param);
//        result.blockingForEach(message -> handleGenerationResult(message));
        return result;
    }

    private static GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
//                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                .model("qwen-turbo-latest")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
    }

    public static void main(String[] args) {
        try {
            Flowable<GenerationResult> result = streamCallWithMessage("中国国土面积多大");
            result.blockingForEach(message -> handleGenerationResult(message));
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            System.err.println("An exception occurred: {}" + e.getMessage());
        }
        System.exit(0);
    }
}
